package com.haredot.interceptor;

import com.haredot.annotation.IpRate;
import com.haredot.interceptor.properties.IpRateProperties;
import com.haredot.interceptor.exception.IpRateException;
import com.haredot.utils.ClientIPUtil;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.List;

/**
 * IP 限流 默认 每秒最多允许访问5次
 */
public class IpRateInterceptor implements HandlerInterceptor {

    private IpRateProperties properties ;

    private StringRedisTemplate stringRedisTemplate ;

    public  IpRateInterceptor(IpRateProperties properties, StringRedisTemplate stringRedisTemplate) {
        this.properties = properties ;
        this.stringRedisTemplate = stringRedisTemplate ;
    }
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {

        // 判断请求是否是 控制器映射的方法， 如果不是，则不需要做限流操作
        if (!(handler instanceof HandlerMethod) || !properties.isEnable()) {
            return true ;
        }

        HandlerMethod handlerMethod = (HandlerMethod) handler;
        // 获取 IP 地址
        String ipAddr = ClientIPUtil.getClientIP(request) ;

        // 怎么限流 ?
        // 获取过期时间，单位毫秒
        long millisExpired = properties.getUnit().toMillis(properties.getTime());
        // 获取在规定时间允许访问的最大次数
        long maxCount = properties.getMaxCount() ;

        // 获取 handlerMethod 对应的 方法
        IpRate ipRate = handlerMethod.getMethodAnnotation(IpRate.class);

        if (ipRate != null) {
            if (ipRate.time() != -1) {
                millisExpired = ipRate.unit().toMillis(ipRate.time());
            }
            if (ipRate.maxCount() != -1) {
                maxCount = ipRate.maxCount() ;
            }
        }
        // 编写一个 Lua 脚本， 进行 IP 限流
        String luaScript = """
                local ipaddr = KEYS[1] ; 
                local count = redis.call('get', ipaddr) ;  
                if not count then 
                   count = '0' ;
                   redis.call('set', ipaddr , count) ; 
                   redis.call('pexpire' , ipaddr, ARGV[1])
                end 
                
                if tonumber(count) >= tonumber(ARGV[2])  then
                    return false ;
                else
                    redis.call('incr', ipaddr) ;
                    return true ;        
                end
                """ ;

        RedisScript<Boolean> redisScript = new DefaultRedisScript<>(luaScript, boolean.class);

        String requestURL = request.getRequestURL().toString();
        // 限流 Key
        String key = ipAddr + "::" + Base64.getEncoder().encodeToString(requestURL.getBytes(StandardCharsets.UTF_8));

        // 执行 Lua 脚本
        boolean luaResult = stringRedisTemplate.execute(redisScript, List.of(key), millisExpired + "", maxCount + "") ;

        if (!luaResult) {
            throw new IpRateException("访问频率过高") ;
        }

        return luaResult ;
    }
}
